#include "CFEMNonlinear.h"
#include <time.h>
#include <fstream>
#include <vtkUnstructuredGrid.h>
#include <vtkSmartPointer.h>
#include <vtkCellArray.h>
#include <vtkTriangle.h>
#include <vtkQuad.h>
#include <vtkQuadraticTriangle.h>
#include <vtkBiQuadraticQuad.h>
#include <vtkXMLUnstructuredGridWriter.h>
#include <vtkArray.h>
#include <vtkDoubleArray.h>
#include <vtkPointData.h>

///@details The coefficient function is the finite element function, and this function has been implemented in CFEMNonlinear::localFESol.
///So in this function we don't need 
SpMat CFEMNonlinear::assembleMatrix2DNonlinear(
	Rsv &uh, double t, CBasisFunction &basis_uh, const int &basis_der_x_uh, const int &basis_der_y_uh, vector<CElement>&Tb,
	CBasisFunction &basis_trial, const int &basis_trial_der_x, const int &basis_trial_der_y, vector<CElement>&Tb_trial,
	CBasisFunction &basis_test, const int &basis_test_der_x, const int &basis_test_der_y, vector<CElement>&Tb_test, const pair<int, int> &matrix_size)
{
	clock_t start, end;
	cout << "\nAssemble matrix" << endl;
	start = clock();

	int estimation_of_entries = m_Mesh.getTbTrial()[0].get_ele_node_num()*m_Mesh.getTbTrial().size(); //estimation of the non-zero values
	vector<Tri> tripletList;//a triple list for the matrix to improve the efficiency of sparse matrix manipulation
	tripletList.reserve(estimation_of_entries);

	for (unsigned int n = 0; n < m_Mesh.getTbTrial().size(); ++n)
	{
		//vertices (coordinates) of nth element(mesh nodes not finite element nodes)	
		vector<CPoint> vertices;	
		for (int i = 0; i < m_Mesh.getElements()[n].get_ele_node_num(); ++i)
		{
			vertices.emplace_back(m_Mesh.getPoints()[m_Mesh.getElements()[n].getIds()[i]]);
		}

		//get the solutions on element n
		vector<double> uh_vector;
		for (unsigned int i = 0; i < Tb[n].getIds().size(); ++i)
		{
			uh_vector.emplace_back(uh(Tb[n].getIds()[i]));
		}	

		vector<double> Gauss_weights;
		vector<CPoint> Gauss_nodes;

		//generate the Gauss integration points	
		generateGaussLocal2D(vertices, m_Mesh.getElements()[n].get_ele_type(), m_GaussType, Gauss_nodes, Gauss_weights);


		for (unsigned int alpha = 0; alpha < Tb_trial[n].getIds().size(); ++alpha)
			for (unsigned int beta = 0; beta < Tb_test[n].getIds().size(); ++beta)
			{
				double int_value = GaussQuad2DCoeTrialTest(
					t, uh_vector, basis_uh,basis_der_x_uh,basis_der_y_uh,Gauss_weights,Gauss_nodes, vertices,
					basis_trial, alpha, basis_trial_der_x, basis_trial_der_y,
					basis_test, beta, basis_test_der_x, basis_test_der_y);

				int i = Tb_test[n].getIds()[beta];
				int j = Tb_trial[n].getIds()[alpha];
				tripletList.emplace_back(Tri(i, j, int_value));
			}
	}
	SpMat A(matrix_size.first, matrix_size.second);
	A.setFromTriplets(tripletList.begin(), tripletList.end());
	end = clock();
	cout << "Time consumption of matrix assemble:\t" << (double)(end - start) / CLOCKS_PER_SEC << " sec" << endl;
	return A;
}

Rsv CFEMNonlinear::assembleVector2DNonlinear(double t,//time for the calculation of uh1 and uh2
	Rsv &uh1,  CBasisFunction &basis_uh1, const int &basis_der_x_uh1, const int &basis_der_y_uh1, vector<CElement> &Tbuh1, //parameters for the calculation of uh1
	Rsv &uh2,  CBasisFunction &basis_uh2, const int &basis_der_x_uh2, const int &basis_der_y_uh2, vector<CElement> &Tbuh2, //parameters for the calculation of uh2
	CBasisFunction & basis_test, const int &basis_der_x_test, const int &basis_der_y_test, vector<CElement> &Tb_test)
{
	clock_t start, end;
	cout << "\nAssemble right side vector" << endl;
	start = clock();
	Rsv b = Rsv(m_Mesh.getPbTest().size());

	b.setZero();

	for (unsigned int n = 0; n < m_Mesh.getTbTest().size(); ++n)
	{
		vector<CPoint> vertices;//vertices of nth element
		vector<double> Gauss_weights;
		vector<CPoint> Gauss_nodes;

		for (int i = 0; i < m_Mesh.getElements()[n].get_ele_node_num(); ++i)
		{
			vertices.emplace_back(m_Mesh.getPoints()[m_Mesh.getElements()[n].getIds()[i]]);
		}
		//generate the Gauss integration points
		generateGaussLocal2D(vertices, m_Mesh.getElements()[n].get_ele_type(), m_GaussType, Gauss_nodes, Gauss_weights);

		//get the solutions on element n
		vector<double> uh1_vector,uh2_vector;
		for (unsigned int i = 0; i < Tbuh1[n].getIds().size(); ++i)
		{
			uh1_vector.emplace_back(uh1(Tbuh1[n].getIds()[i]));
		}
		for (unsigned int i = 0; i < Tbuh2[n].getIds().size(); ++i)
		{
			uh2_vector.emplace_back(uh2(Tbuh2[n].getIds()[i]));
		}

		for (size_t beta = 0; beta < Tb_test[n].getIds().size(); ++beta)
		{
			double int_value = GaussQuad2DLoadTrial(t,
				uh1_vector,basis_uh1,basis_der_x_uh1,basis_der_y_uh1,//u1
				uh2_vector,basis_uh2,basis_der_x_uh2,basis_der_y_uh2,//u2
				Gauss_weights, Gauss_nodes, vertices, basis_test, beta, basis_der_x_test, basis_der_y_test);

			b[Tb_test[n].getIds()[beta]] += int_value;
		}
	}
	end = clock();

	cout << "Time consumption of right side vector assemble:\t" << (double)(end - start) / CLOCKS_PER_SEC << " sec" << endl << endl;

	return b;
}

void CFEMNonlinear::export2VTKXml(const string &filename, const string &dataArrayName, const double &time,const Rsv & x, vector<CPoint>& Pb, vector<CElement>& Tb, const int &vec_size)
{
	cout << "\nExport results to vtk with xml format visualized with Paraview" << endl;
	/*ofstream f(filename);
	f << "<? xml version = \"1.0\" encoding = \"UTF - 8\" ? >" << endl;
	f << "<VTKFile type = \"UnstructuredGrid\" version = \"0.1\" byte_order = \"LittleEndian\">" << endl;
	f << "<UnstructuredGrid>" << endl;
	f << "<Piece NumberOfPoints = \"" << Pb.size() << "\" NumberOfCells = \"" << Tb.size() << "\">" << endl;
	f << "<PointData>" << endl;
	
	f << "<DataArray type = \"Float64\" Name = \"s_fem\" Format = \"ascii\">" << endl;*/

	vtkSmartPointer<vtkPoints> points =
		vtkSmartPointer<vtkPoints>::New();
	for (unsigned int i=0;i<Pb.size();++i)
	{
		points->InsertNextPoint(Pb[i].getX(), Pb[i].getY(), Pb[i].getZ());
	}

	vtkSmartPointer<vtkCellArray> cellArray =
		vtkSmartPointer<vtkCellArray>::New();

	std::vector<int> cellTypes;
	
	
	for (unsigned int i=0;i<Tb.size();++i)
	{
		
		if (ELE_3==Tb[i].get_ele_type())
		{
			vtkSmartPointer <vtkTriangle> cell= vtkSmartPointer <vtkTriangle>::New();
			for (unsigned int j=0;j<Tb[i].getIds().size();++j)
			{
				cell->GetPointIds()->SetId(j, Tb[i].getIds()[j]);
			}
			cellArray->InsertNextCell(cell);
			cellTypes.emplace_back(VTK_TRIANGLE);
		}
		if (ELE_4 == Tb[i].get_ele_type())
		{
			vtkSmartPointer <vtkQuad> cell= vtkSmartPointer <vtkQuad>::New();
			for (unsigned int j = 0; j < Tb[i].getIds().size(); ++j)
			{
				cell->GetPointIds()->SetId(j, Tb[i].getIds()[j]);
			}
			cellArray->InsertNextCell(cell);
			cellTypes.emplace_back(VTK_QUAD);
		}
		if (ELE_6 == Tb[i].get_ele_type())
		{
			vtkSmartPointer<vtkQuadraticTriangle> cell= vtkSmartPointer<vtkQuadraticTriangle>::New();
			for (unsigned int j = 0; j < Tb[i].getIds().size(); j=j+2)
			{
				cell->GetPointIds()->SetId(j/2, Tb[i].getIds()[j]);
			}
			for (unsigned int j = 1; j < Tb[i].getIds().size(); j = j + 2)
			{
				cell->GetPointIds()->SetId(3+(j-1)/2, Tb[i].getIds()[j]);
			}
			cellArray->InsertNextCell(cell);
			cellTypes.emplace_back(VTK_QUADRATIC_TRIANGLE);
		}
		if (ELE_9 == Tb[i].get_ele_type())
		{
			vtkSmartPointer<vtkBiQuadraticQuad> cell= vtkSmartPointer<vtkBiQuadraticQuad>::New();
			for (unsigned int j = 0; j < Tb[i].getIds().size(); j = j + 2)
			{
				cell->GetPointIds()->SetId(j/2, Tb[i].getIds()[j]);
			}
			for (unsigned int j = 1; j < Tb[i].getIds().size(); j = j + 2)
			{
				cell->GetPointIds()->SetId((j-1)/2 + 4, Tb[i].getIds()[j]);
			}
			cell->GetPointIds()->SetId(8, Tb[i].getIds()[8]);
			cellArray->InsertNextCell(cell);
			cellTypes.emplace_back(VTK_BIQUADRATIC_QUAD);
		}
	}

	vtkSmartPointer<vtkUnstructuredGrid> unstructuredGrid =
		vtkSmartPointer<vtkUnstructuredGrid>::New();
	unstructuredGrid->SetPoints(points);
	unstructuredGrid->SetCells(&(cellTypes[0]), cellArray);

	//set points data
	if (vec_size == 1)
	{
		vtkSmartPointer <vtkDoubleArray> scalar_value=
			vtkSmartPointer <vtkDoubleArray>::New();
		scalar_value->SetNumberOfValues(Pb.size());
		scalar_value->SetName(dataArrayName.c_str());

		for (unsigned int i = 0; i < Pb.size(); ++i)
		{
			scalar_value->SetValue(i, x(i));
		}
		unstructuredGrid->GetPointData()->SetScalars(scalar_value);
	}
	else
	{
		//set vertex vector
		vtkSmartPointer<vtkDoubleArray> vector_value =
			vtkSmartPointer<vtkDoubleArray>::New();
		vector_value->SetNumberOfComponents(3);
		vector_value->SetNumberOfTuples(Pb.size());
		vector_value->SetName(dataArrayName.c_str());
		//construct the normal vectors
		for (int i=0;i<Pb.size();++i)
		{
			double data[3];
			int j;
			for (j=0;j<vec_size;++j)
			{
				data[j] = x(i + j * Pb.size());
			}
			for (int k = j; k < 3 ; ++k)
			{
				data[k]=0;
			}
			vector_value->SetTuple(i, data);
		}			
		unstructuredGrid->GetPointData()->SetVectors(vector_value);
	}
	//set the time
	vtkDoubleArray *t = vtkDoubleArray::New();
	t->SetName("TIME");
	t->SetNumberOfTuples(1);
	t->SetTuple1(0, time);
	unstructuredGrid->GetFieldData()->AddArray(t);

	// Write file
	vtkSmartPointer<vtkXMLUnstructuredGridWriter> writer =
		vtkSmartPointer<vtkXMLUnstructuredGridWriter>::New();
	writer->SetFileName(filename.c_str());
#if VTK_MAJOR_VERSION <= 5
	writer->SetInput(unstructuredGrid);
#else
	writer->SetInputData(unstructuredGrid);
#endif
	writer->Write();


}

void CFEMNonlinear::export2VTK(string filename, const Rsv & x, vector<CPoint>& P, vector<CElement>& T, int num_vec)
{
	cout << "\nExport results to vtk visualized with Paraview" << endl;
	map<ELE_TYPE, int> type_ele_vtk = {
		{ELE_3,5},
		{ELE_4,9},
		{ELE_6,22},
		{ELE_9,28}
	};
	ofstream f(filename);
	f << "# vtk DataFile Version 3.0" << endl;
	f << "Result" << endl;
	f << "ASCII" << endl;
	f << "DATASET UNSTRUCTURED_GRID" << endl;

	f << "POINTS \t" << P.size() << "\tdouble" << endl;
	for (unsigned int i = 0; i < P.size(); ++i)
	{
		f << P[i].getX() << "\t" << P[i].getY() << "\t" << P[i].getZ() << endl;
	}
	f << "CELLS \t " << T.size() << "\t" << (T[0].getIds().size() + 1)* T.size() << endl;
	for (unsigned int i = 0; i < T.size(); ++i)
	{

		f << T[i].getIds().size() << "\t";

		if (T[i].get_ele_type() == ELE_3 || T[i].get_ele_type() == ELE_4)
		{
			for (unsigned int j = 0; j < T[i].getIds().size(); ++j)
			{
				f << T[i].getIds()[j] << "\t";
			}
		}
		else if (T[i].get_ele_type() == ELE_6)
		{
			for (unsigned int j = 0; j < T[i].getIds().size(); j += 2)
			{
				f << T[i].getIds()[j] << "\t";
			}
			for (unsigned int j = 1; j < T[i].getIds().size(); j += 2)
			{
				f << T[i].getIds()[j] << "\t";
			}
		}
		else if (T[i].get_ele_type() == ELE_9)
		{
			for (unsigned int j = 0; j < T[i].getIds().size() - 1; j += 2)
			{
				f << T[i].getIds()[j] << "\t";
			}
			for (unsigned int j = 1; j < T[i].getIds().size() - 1; j += 2)
			{
				f << T[i].getIds()[j] << "\t";
			}
			f << T[i].getIds()[8] << "\t";
		}
		f << endl;

	}
	f << "CELL_TYPES\t" << T.size() << endl;

	for (unsigned int i = 0; i < T.size(); ++i)
	{
		f << type_ele_vtk[T[i].get_ele_type()] << endl;
	}

	f << "POINT_DATA\t" << P.size() << endl;
	if (num_vec==1)
	{
		f << "SCALARS s_fem double 1" << endl;
		f << "LOOKUP_TABLE default" << endl;
		for (unsigned int i = 0; i < P.size(); ++i)
		{
			f << x(i) << endl;
		}
		f.close();
	}
	else
	{
		f << "VECTORS s_fem double" << endl;
		for (unsigned int i=0;i<P.size();++i)
		{
			int j;
			for (j=0;j<num_vec;++j)
			{
				f << x(i + j * P.size()) << "\t";
			}
			for (int k=0;k<3-j;++k)
			{
				f << 0;
			}
			f << endl;
		}
		f.close();
	}
	
	
	
}

double CFEMNonlinear::localFESol(const vector<double>& uh_vector, const CPoint & x, const vector<CPoint>& vertices, 
	CBasisFunction & basis_fun, const int & basis_der_x, const int & basis_der_y)
{
	double result = 0;
	for (int k = 0; k < basis_fun.getNumberBasis(); ++k)
	{
		result += uh_vector[k] * basis_fun.FE_basis_local_func_2D(x, vertices, basis_fun.getType(), k, basis_der_x, basis_der_y);

	}
	return result;
}

double CFEMNonlinear::GaussQuad2DCoeTrialTest(double t, const vector<double>&uh_vector, CBasisFunction &basis_fun_uh, const int &basis_der_x_uh, const int &basis_der_y_uh,//parameters needed for the calculation of finite element
	const vector<double>& Gauss_weights, const vector<CPoint>& Gauss_nodes, const vector<CPoint>& vertices, 
	CBasisFunction & basis_trial, const int & basis_index_trial, const int & basis_der_x_trial, const int & basis_der_y_trial, 
	CBasisFunction & basis_test, const int & basis_index_test, const int & basis_der_x_test, const int & basis_der_y_test)
{
	int Gpn = Gauss_weights.size();
	double int_value = 0;

	for (int k = 0; k < Gpn; ++k)
	{
		int_value += Gauss_weights[k] * localFESol(uh_vector, Gauss_nodes[k], vertices, basis_fun_uh, basis_der_x_uh, basis_der_y_uh)
			*basis_trial.FE_basis_local_func_2D(Gauss_nodes[k], vertices, basis_trial.getType(), basis_index_trial, basis_der_x_trial, basis_der_y_trial)
			*basis_test.FE_basis_local_func_2D(Gauss_nodes[k], vertices, basis_test.getType(), basis_index_test, basis_der_x_test, basis_der_y_test);
	}
	return int_value;
}
double CFEMNonlinear::GaussQuad2DLoadTrial(double t,//time for the calculation of uh1 and uh2
	const vector<double> uh1_vector,  CBasisFunction &basis_uh1, const int &basis_der_x_uh1, const int &basis_der_y_uh1, //parameters for the calculation of uh1
	const vector<double> uh2_vector,  CBasisFunction &basis_uh2, const int &basis_der_x_uh2, const int &basis_der_y_uh2,  //parameters for the calculation of uh2	
	const vector<double> &Gauss_weights, const vector<CPoint> &Gauss_nodes, const vector<CPoint> &vertices,
	CBasisFunction &basis_test, const int &basis_index_test, const int &basis_der_x_test, const int &basis_der_y_test)
{
	int Gpn = Gauss_weights.size();
	double int_value = 0;

	for (int k = 0; k < Gpn; ++k)
	{
		int_value += 
			Gauss_weights[k]
			*localFESol(uh1_vector, Gauss_nodes[k], vertices, basis_uh1, basis_der_x_uh1, basis_der_y_uh1)
			*localFESol(uh2_vector, Gauss_nodes[k], vertices, basis_uh2, basis_der_x_uh2, basis_der_y_uh2)
			*basis_test.FE_basis_local_func_2D(Gauss_nodes[k], vertices, basis_test.getType(), basis_index_test, basis_der_x_test, basis_der_y_test);
	}
	return int_value;

}
